package consistent_hash

import (
	"hash/crc32"
	"sort"
	"strconv"
)

// ConsistentHash 定义一致性哈希结构
type ConsistentHash struct {
	hashRing     []uint32          // 哈希环
	virtualNodes map[uint32]string // 虚拟节点到真实节点的映射
	replicas     int               // 每个真实节点的虚拟节点数
	nodes        map[string]bool
	hashFunction func(data []byte) uint32
}

// NewConsistentHash 创建一个新的 ConsistentHash 实例
func NewConsistentHash(replicas int) *ConsistentHash {
	return &ConsistentHash{
		hashRing:     make([]uint32, 0),
		virtualNodes: make(map[uint32]string),
		replicas:     replicas,
		nodes:        make(map[string]bool),
		hashFunction: crc32.ChecksumIEEE,
	}
}

// AddNode 添加一个节点到一致性哈希环中
func (ch *ConsistentHash) AddNode(node string) {
	// 节点已存在，无需重复添加
	if ch.nodes[node] {
		return
	}
	ch.nodes[node] = true

	// 创建一个临时切片来存储新的虚拟节点
	newVirtualNodes := make([]uint32, 0, ch.replicas)

	for i := 0; i < ch.replicas; i++ {
		virtualNodeKey := node + "#" + strconv.Itoa(i)
		hash := ch.hashFunction([]byte(virtualNodeKey))
		newVirtualNodes = append(newVirtualNodes, hash)
		ch.virtualNodes[hash] = node
	}

	// 对新虚拟节点进行排序，确保均匀分布
	sort.Slice(newVirtualNodes, func(i, j int) bool {
		return newVirtualNodes[i] < newVirtualNodes[j]
	})

	// 将排序后的虚拟节点插入到哈希环中，保持哈希环有序
	for _, hash := range newVirtualNodes {
		index := sort.Search(len(ch.hashRing), func(i int) bool {
			return ch.hashRing[i] >= hash
		})
		ch.hashRing = append(ch.hashRing[:index], append([]uint32{hash}, ch.hashRing[index:]...)...)
	}
}

// RemoveNode 从一致性哈希环中移除一个节点
func (ch *ConsistentHash) RemoveNode(node string) {
	if !ch.nodes[node] {
		return // 节点不存在，无需移除
	}
	delete(ch.nodes, node)
	for i := 0; i < ch.replicas; i++ {
		virtualNodeKey := node + "#" + strconv.Itoa(i)
		hash := ch.hashFunction([]byte(virtualNodeKey))
		delete(ch.virtualNodes, hash)
		// 从哈希环中移除对应的虚拟节点
		index := sort.Search(len(ch.hashRing), func(i int) bool {
			return ch.hashRing[i] >= hash
		})
		if index < len(ch.hashRing) && ch.hashRing[index] == hash {
			ch.hashRing = append(ch.hashRing[:index], ch.hashRing[index+1:]...)
		}
	}
}

// GetNode 根据 key 获取对应的真实节点
func (ch *ConsistentHash) GetNode(key string) string {
	if len(ch.hashRing) == 0 {
		return ""
	}

	hash := ch.hashFunction([]byte(key))
	index := sort.Search(len(ch.hashRing), func(i int) bool {
		return ch.hashRing[i] >= hash
	}) % len(ch.hashRing)

	return ch.virtualNodes[ch.hashRing[index]]
}
